# -*- coding: utf-8 -*-
"""
@Time ： 2020-11-19 9:30
@Auth ： lixin
@Description：数值型数据分类算法集合训练类

"""
from sklearn import linear_model

from algo.Algo_interface import Algo_interface
from sklearn import linear_model
from sklearn import metrics
from sklearn.ensemble import RandomForestClassifier
from sklearn import svm
from sklearn.neighbors import KNeighborsClassifier
# import api as runs
import numpy as np
import lightgbm as lgb
from sklearn import neural_network
from xgboost.sklearn import XGBClassifier
from sklearn.ensemble import GradientBoostingClassifier
import json
from sklearn.naive_bayes import MultinomialNB

class SklearnModel(Algo_interface):

    def __init__(self, model_type):
        self.task_type = model_type
        self.model = None
        # return self.model

    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        return self.model

    def train(self,data):
        x_train, y_train=data
        self.model.fit(x_train, y_train)
        return 1

    def predict(self,data):
        y_pred = self.model.predict(data)
        return y_pred